import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
from torchvision import datasets, transforms
import numpy as np
from PIL import Image
import random
import matplotlib.pyplot as plt
from tqdm import tqdm
import os
import cv2
from skimage.metrics import peak_signal_noise_ratio as psnr
from skimage.metrics import structural_similarity as ssim
import argparse

# 设置设备
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print("Using device:", device)

# 解析命令行参数，指定目标epoch
parser = argparse.ArgumentParser(description='DnCNN denoising with traditional filters and full metrics')
parser.add_argument('--target_epoch', type=int, required=True, help='Target epoch number for training or inference')
parser.add_argument('--test_image', type=str, required=True, help='Path to test image')
parser.add_argument('--usb_path', type=str, required=True, help='Path to USB drive for saving results')
args = parser.parse_args()

# 定义保存路径
local_save_dir = os.path.join(os.getcwd(), 'results')  # 实验室电脑本地路径
usb_save_dir = os.path.join(args.usb_path, 'denoising_results')  # U盘路径

# 创建保存目录
os.makedirs(local_save_dir, exist_ok=True)
os.makedirs(usb_save_dir, exist_ok=True)

# 下载CIFAR10数据集，转灰度
transform_gray = transforms.Compose([
    transforms.ToTensor(),
    transforms.Lambda(lambda x: x.mean(dim=0, keepdim=True))  # RGB转灰度通道
])

trainset = datasets.CIFAR10(root='./data', train=True, download=True, transform=transform_gray)

class DnCNNDataset(Dataset):
    def __init__(self, base_dataset, patch_size=32, sigma=25):
        self.base_dataset = base_dataset
        self.patch_size = patch_size
        self.sigma = sigma / 255.0  # 归一化

    def __len__(self):
        return len(self.base_dataset)

    def __getitem__(self, idx):
        img, _ = self.base_dataset[idx]  # Tensor 1xHxW
        c, h, w = img.shape
        top = random.randint(0, h - self.patch_size)
        left = random.randint(0, w - self.patch_size)
        patch = img[:, top:top+self.patch_size, left:left+self.patch_size]
        noise = torch.randn_like(patch) * self.sigma
        noisy_patch = patch + noise
        noisy_patch = torch.clamp(noisy_patch, 0., 1.)

        return noisy_patch, patch

# DnCNN模型改为单通道输入输出
class DnCNN(nn.Module):
    def __init__(self, channels=1, num_layers=17):
        super(DnCNN, self).__init__()
        layers = [nn.Conv2d(channels, 64, kernel_size=3, padding=1), nn.ReLU(inplace=True)]
        for _ in range(num_layers - 2):
            layers += [
                nn.Conv2d(64, 64, kernel_size=3, padding=1),
                nn.BatchNorm2d(64),
                nn.ReLU(inplace=True)
            ]
        layers.append(nn.Conv2d(64, channels, kernel_size=3, padding=1))
        self.dncnn = nn.Sequential(*layers)

    def forward(self, x):
        return x - self.dncnn(x)

# 初始化模型
model = DnCNN(channels=1).to(device)

# 损失函数和优化器
criterion = nn.MSELoss()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

# 准备训练数据
train_dataset = DnCNNDataset(trainset, patch_size=32, sigma=25)
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)

# 检查是否存在目标epoch的模型
def get_model_path(epoch):
    return f"dncnn_gray_epoch_{epoch}.pth"

target_epoch = args.target_epoch
model_filename = get_model_path(target_epoch)
local_model_path = os.path.join(local_save_dir, model_filename)
usb_model_path = os.path.join(usb_save_dir, model_filename)

# 检查本地或U盘是否已有该模型
model_exists = False
if os.path.exists(local_model_path):
    model.load_state_dict(torch.load(local_model_path, map_location=device))
    model_exists = True
elif os.path.exists(usb_model_path):
    model.load_state_dict(torch.load(usb_model_path, map_location=device))
    model_exists = True

# 如果模型不存在，则从头重新训练（删除了继续训练的逻辑）
if not model_exists:
    print(f"模型 epoch {target_epoch} 不存在，从头开始训练...")

    # 直接从头训练到目标epoch（无继续训练逻辑）
    model.train()
    for epoch in range(target_epoch):
        running_loss = 0
        for noisy, clean in tqdm(train_loader, desc=f"Epoch {epoch+1}/{target_epoch}"):
            noisy, clean = noisy.to(device), clean.to(device)
            output = model(noisy)
            loss = criterion(output, clean)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            running_loss += loss.item()
        print(f"Epoch {epoch+1} average loss: {running_loss / len(train_loader):.6f}")

    # 保存模型到本地和U盘
    torch.save(model.state_dict(), local_model_path)
    torch.save(model.state_dict(), usb_model_path)
    print(f"模型已保存到本地: {local_model_path}")
    print(f"模型已保存到U盘: {usb_model_path}")
else:
    print(f"已加载现有模型 epoch {target_epoch}")

# -------------------------- 测试图像处理（核心扩展部分）--------------------------
test_image_path = args.test_image
if os.path.exists(test_image_path):
    # 1. 读取并预处理原始图像
    img = Image.open(test_image_path).convert('L')  # 转灰度
    img = img.resize((128, 128))  # 统一尺寸
    img_np = np.array(img).astype(np.float32) / 255.0  # 归一化到0-1（用于计算指标）
    img_uint8 = np.array(img).astype(np.uint8)  # 原始uint8格式（用于滤波操作）

    # 2. 生成加噪图像
    sigma = 25 / 255.0
    noise = torch.randn_like(torch.from_numpy(img_np).unsqueeze(0).unsqueeze(0)) * sigma
    noisy_tensor = torch.from_numpy(img_np).unsqueeze(0).unsqueeze(0) + noise
    noisy_tensor = torch.clamp(noisy_tensor, 0., 1.)
    noisy_np = noisy_tensor.squeeze().cpu().numpy()  # 0-1格式（用于计算指标）
    noisy_uint8 = (noisy_np * 255).astype(np.uint8)  # uint8格式（用于滤波操作）

    # 3. DnCNN模型去噪
    model.eval()
    with torch.no_grad():
        dncnn_denoised = model(noisy_tensor.to(device)).clamp(0, 1)
    dncnn_denoised_np = dncnn_denoised.squeeze().cpu().numpy()  # 0-1格式
    dncnn_denoised_uint8 = (dncnn_denoised_np * 255).astype(np.uint8)  # uint8格式

    # 4. 传统滤波去噪（新增）
    # 4.1 中值滤波（常用去噪，核大小3x3）
    median_filtered = cv2.medianBlur(noisy_uint8, ksize=3)
    median_filtered_np = median_filtered.astype(np.float32) / 255.0  # 归一化用于指标计算

    # 4.2 高斯滤波（平滑去噪，核5x5，sigmaX=1.5）
    gaussian_filtered = cv2.GaussianBlur(noisy_uint8, ksize=(5, 5), sigmaX=1.5)
    gaussian_filtered_np = gaussian_filtered.astype(np.float32) / 255.0

    # 4.3 双边滤波（保边去噪，颜色标准差10，空间标准差15）
    bilateral_filtered = cv2.bilateralFilter(noisy_uint8, d=9, sigmaColor=10, sigmaSpace=15)
    bilateral_filtered_np = bilateral_filtered.astype(np.float32) / 255.0

    # 5. 计算所有需要的PSNR和SSIM（核心补充）
    metrics = {}
    # 5.1 原始图 vs 加噪图
    metrics['original_vs_noisy'] = {
        'psnr': psnr(img_np, noisy_np, data_range=1.0),
        'ssim': ssim(img_np, noisy_np, data_range=1.0)
    }
    # 5.2 原始图 vs DnCNN去噪图
    metrics['original_vs_dncnn'] = {
        'psnr': psnr(img_np, dncnn_denoised_np, data_range=1.0),
        'ssim': ssim(img_np, dncnn_denoised_np, data_range=1.0)
    }
    # 5.3 原始图 vs 中值滤波图
    metrics['original_vs_median'] = {
        'psnr': psnr(img_np, median_filtered_np, data_range=1.0),
        'ssim': ssim(img_np, median_filtered_np, data_range=1.0)
    }
    # 5.4 原始图 vs 高斯滤波图
    metrics['original_vs_gaussian'] = {
        'psnr': psnr(img_np, gaussian_filtered_np, data_range=1.0),
        'ssim': ssim(img_np, gaussian_filtered_np, data_range=1.0)
    }
    # 5.5 原始图 vs 双边滤波图
    metrics['original_vs_bilateral'] = {
        'psnr': psnr(img_np, bilateral_filtered_np, data_range=1.0),
        'ssim': ssim(img_np, bilateral_filtered_np, data_range=1.0)
    }

    # 打印所有指标
    print("\n" + "="*50)
    print(f"所有PSNR/SSIM指标（epoch={target_epoch}）")
    print("="*50)
    for pair, score in metrics.items():
        print(f"{pair:<20} | PSNR: {score['psnr']:.2f} dB | SSIM: {score['ssim']:.4f}")
    print("="*50 + "\n")

    # 6. 保存所有指标到文件（本地+U盘）
    metrics_filename = f"full_metrics_epoch_{target_epoch}.txt"
    metrics_content = f"Epoch: {target_epoch}\n"
    metrics_content += "="*50 + "\n"
    metrics_content += "指标说明：所有结果均为「对比图 vs 原始图」\n"
    metrics_content += "="*50 + "\n"
    for pair, score in metrics.items():
        metrics_content += f"{pair:<20} | PSNR: {score['psnr']:.2f} dB | SSIM: {score['ssim']:.4f}\n"

    # 保存到本地
    with open(os.path.join(local_save_dir, metrics_filename), 'w') as f:
        f.write(metrics_content)
    # 保存到U盘
    with open(os.path.join(usb_save_dir, metrics_filename), 'w') as f:
        f.write(metrics_content)

    # 7. 显示所有图像对比（原始图+加噪图+4种去噪图）
    plt.figure(figsize=(18, 4))  # 扩大画布以容纳7张子图

    # 7.1 原始图
    plt.subplot(1, 6, 1)
    plt.imshow(img_np, cmap='gray')
    plt.title("Original\nImage")
    plt.axis('off')

    # 7.2 加噪图
    plt.subplot(1, 6, 2)
    plt.imshow(noisy_np, cmap='gray')
    plt.title(f"Noisy\n(PSNR: {metrics['original_vs_noisy']['psnr']:.2f} dB)")
    plt.axis('off')

    # 7.3 DnCNN去噪图
    plt.subplot(1, 6, 3)
    plt.imshow(dncnn_denoised_np, cmap='gray')
    plt.title(f"DnCNN Denoised\n(PSNR: {metrics['original_vs_dncnn']['psnr']:.2f} dB)")
    plt.axis('off')

    # 7.4 中值滤波图
    plt.subplot(1, 6, 4)
    plt.imshow(median_filtered_np, cmap='gray')
    plt.title(f"Median Filter\n(PSNR: {metrics['original_vs_median']['psnr']:.2f} dB)")
    plt.axis('off')

    # 7.5 高斯滤波图
    plt.subplot(1, 6, 5)
    plt.imshow(gaussian_filtered_np, cmap='gray')
    plt.title(f"Gaussian Filter\n(PSNR: {metrics['original_vs_gaussian']['psnr']:.2f} dB)")
    plt.axis('off')

    # 7.6 双边滤波图
    plt.subplot(1, 6, 6)
    plt.imshow(bilateral_filtered_np, cmap='gray')
    plt.title(f"Bilateral Filter\n(PSNR: {metrics['original_vs_bilateral']['psnr']:.2f} dB)")
    plt.axis('off')

    plt.tight_layout()
    plt.show()

    # 8. 保存所有图片（本地+U盘）
    # 定义文件名和对应图像
    image_files = {
        f"original_epoch_{target_epoch}.png": img_uint8,
        f"noisy_epoch_{target_epoch}.png": noisy_uint8,
        f"dncnn_denoised_epoch_{target_epoch}.png": dncnn_denoised_uint8,
        f"median_filtered_epoch_{target_epoch}.png": median_filtered,
        f"gaussian_filtered_epoch_{target_epoch}.png":gaussian_filtered_epoch,
        f"bilateral_filtered_epoch_{target_epoch}.png":bilateral_filtered
    }
    # 批量保存到本地和U盘
    for filename, img_data in image_files.items():
        # 保存到本地
        local_path = os.path.join(local_save_dir, filename)
        Image.fromarray(img_data).save(local_path)
        # 保存到U盘
        usb_path = os.path.join(usb_save_dir, filename)
        Image.fromarray(img_data).save(usb_path)
        print(f"{filename:<30} 已保存到本地: {local_path}")
        print(f"{filename:<30} 已保存到U盘: {usb_path}")

    print(f"\n所有结果（图片+指标）已保存完成！")
else:
    print(f"未找到测试图像: {test_image_path}")